// standard headers
#include <istream>
#include <vector>
#include <string>
#include <deque>
#include <fstream>
#include <sstream>
#include <iostream>
#include <algorithm>
#include <stdexcept>
#include <map>
#include <cstdlib>

// boost headers
#include <boost/format.hpp>
#include <boost/foreach.hpp>
#include <boost/lexical_cast.hpp>
#include <boost/filesystem/path.hpp>

// components headers
#include <tree.h>
#include <tree_ops.h>
#include <portable_timer.h>
#include <progress_indicator.h>
#include <print_error.h>
#define DEBUG_LEVEL 6
#include <debug_func.h>
#include <matrix.h>
#include <matrix_io.h>
#include "get_arg.h"




using boost::format;
using boost::lexical_cast;
using boost::filesystem::path;
using std::deque;
using std::vector;
using std::ostringstream;

using std::string;
using std::fill;
using std::ofstream;
using std::cout;
using std::cin;
using std::cerr;
using std::ifstream;
using std::map;
using std::pair;
using std::istream;
using std::max;
using std::min;
using namespace bcpp_matrix;
using namespace bcpp_tree_ops;


//________________________________________________________________________________________

//	do_gene_to_species_map

//		Contains matrix for:
// 		1) transitions
// 			determining the corresponding type of internal nodes
//				on a gene tree when mapped to a species tree
// 		2) deletions
// 			the mapping function in terms of missing/deleted genes
// 		3) duplications
// 			the mapping function in terms of inferred gene duplications

//________________________________________________________________________________________
struct do_gene_to_species_map
{
										// number of nodes in species tree
	unsigned							cnt_tree_nodes;
	t_dense_matrix<unsigned>			duplications;
	t_dense_matrix<unsigned>			deletions;
	t_dense_matrix<unsigned>			transitions;
	t_dense_matrix<unsigned>			speciation_node;
										// pretty names for Last Common Ancestors
	vector<string>						node_names;
										// the ID of the nodes (in sorted order)
	vector<unsigned>					node_identifier;
										// running count to calculate the above
	unsigned							cnt_identifiers;
										// the number of descendent leaves for each node 
	vector<unsigned>					cnt_leaves;
	bcpp_tree::t_tree					tree;
	do_gene_to_species_map(	bcpp_tree::t_tree&				tree_):
			cnt_tree_nodes(tree_.get_node_count()),
			duplications(cnt_tree_nodes + 1, 0),
			deletions(cnt_tree_nodes + 1, 0),
			transitions(cnt_tree_nodes + 1, 65536),
			speciation_node(cnt_tree_nodes + 1, 0),
			node_names(cnt_tree_nodes + 1, string()),
			node_identifier(cnt_tree_nodes, 0),
			cnt_identifiers(0),
			cnt_leaves(cnt_tree_nodes, 0)
	{
		// normalized in alphabetical order to ensure consistency however trees are specified
		tree_.get_normalized_tree(tree);

		// self vs self = duplication and stay at self
		for (unsigned i = 0; i < cnt_tree_nodes; ++i)
		{
			transitions[i][i] = i;
			duplications[i][i] = 1;
		}
		deque<unsigned> child_ids;
		unsigned cnt_leaves_ = 0;
		recurse(tree.get_root(), child_ids, cnt_leaves_);


		// handle "unknown" species
		node_names[cnt_tree_nodes] = "unknown";
		for (unsigned i = 0; i < cnt_tree_nodes + 1; ++i)
		{
			// unknown vs anything is itself
			transitions[i][cnt_tree_nodes] = 
			transitions[cnt_tree_nodes][i] = i;
		}
	}


	
	void recurse(bcpp_tree::t_node* node, deque<unsigned>& child_ids, unsigned& cnt_leaves_);
};

void do_gene_to_species_map::recurse(bcpp_tree::t_node* node, deque<unsigned>& child_ids, unsigned& cnt_leaves_)
{
	debug_func_cerr(5);

	//
	// This is a leaf
	//
	if (node->is_leaf())
	{
		// set the leaf identifier
		unsigned my_id = cnt_identifiers++;
		node_identifier[node->get_UID()] = my_id;
		child_ids.push_back(my_id);

		// set the name and the number of leaves
		node_names[my_id] = node->get_name();
		cnt_leaves[my_id] = 1;
		cnt_leaves_ += 1;
		return;
	}

	// assume binary tree
	assert(node->get_children_count() == 2);


	// get all descendant nodes on the left and right sides
	deque<unsigned> left_child_ids;
	unsigned left_child_leaf_cnt = 0;
	recurse(node->get_child(0), left_child_ids, left_child_leaf_cnt);

	deque<unsigned> rite_child_ids;
	unsigned rite_child_leaf_cnt = 0;
	recurse(node->get_child(1), rite_child_ids, rite_child_leaf_cnt);

	// the number of leaves on this branch
	cnt_leaves_ = left_child_leaf_cnt + rite_child_leaf_cnt;


	// only assign my own id after all children assigned
	unsigned my_id = cnt_identifiers++;
	assert(node->get_UID() < node_identifier.size());
	assert(my_id < cnt_tree_nodes);
	node_identifier[node->get_UID()] = my_id;
	// fancy name!
	node_names[my_id] = node_names[node_identifier[node->get_child(0)->get_UID()]] + "-" +
						node_names[node_identifier[node->get_child(1)->get_UID()]];
	// set node count
	cnt_leaves[my_id]	=	cnt_leaves_;

	// The child_ids of my parent include all my children (and me but later...)
	child_ids = left_child_ids;
	child_ids.insert(child_ids.end(), rite_child_ids.begin(),  rite_child_ids.end());


	// child species meet each other at me
	BOOST_FOREACH(unsigned l, left_child_ids)
		BOOST_FOREACH(unsigned r, rite_child_ids)
		{
			assert(l < cnt_tree_nodes);
			assert(r < cnt_tree_nodes);

			// therefore transitions point to me
			transitions	[l][r] =
			transitions	[r][l] = my_id;

			// this is the point of speciation
			speciation_node	[l][r] =
			speciation_node	[r][l] = 1;
			

			// the number of gene losses is the number of missing leaves
			deletions	[l][r] =
			deletions	[r][l] = cnt_leaves_ - cnt_leaves[l] - cnt_leaves[r];

			// no duplications when child species meets another child species at the
			//	parent node
			
		}


	// when my children species meets me
	BOOST_FOREACH(unsigned c, child_ids)
	{
		assert(c < cnt_tree_nodes);
		// still me but duplication and gene loss
		transitions	[c][my_id]	=
		transitions	[my_id][c]	= my_id;
		duplications[c][my_id]	=
		duplications[my_id][c]	= 1;
		deletions	[c][my_id]	=
		deletions	[my_id][c]	= cnt_leaves_ - cnt_leaves[c];
	}

	// now add my own id to the list of children for my parent
	child_ids.push_back(my_id);
}







//8888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

//	main

//8888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

//
//	print out unsigned matrices
// 
void output_matrix(t_dense_matrix<unsigned>& matrix, vector<string>& names)
{
	t_matrix_io_template<unsigned> output;
	output.swap(matrix, names);
	write_matrix(cout, output,  eMIO_SQUARE, false);
	output.matrix.swap(matrix);
}

int main (int argc, char *argv[])
{
	// we use cout to print progress so flush per call
	//	cout.setf(std::ios_base::unitbuf);

	t_program_args args;
	try
	{

		if (!process_command_line_options(argc, argv, args))
			return 1;

		//
		// read supplied species tree
		//
		unsigned line_cnt = 0;
		bcpp_tree::t_tree species_tree;
		read_tree_from_stream(cin, species_tree, line_cnt);
		if (!cin)
			throw std::runtime_error("No tree supplied.");
	

		// 
		// calculate gene to species mapping functions
		// 
		do_gene_to_species_map species_map(species_tree);

		// 
		// print out the mapping function matrices
		// 
		cout << ">Duplications\n";
		output_matrix(species_map.duplications, species_map.node_names);
		cout << ">Deletions\n";
		output_matrix(species_map.deletions, species_map.node_names);
		cout << ">Transitions\n";
		output_matrix(species_map.transitions, species_map.node_names);
		cout << ">Speciation nodes\n";
		output_matrix(species_map.speciation_node, species_map.node_names);

		//
		// Only for display / debugging
		//
		cout << ">Transition Node names\n";
		for (unsigned i = 0; i < species_map.node_names.size(); ++i)
			cout << "\t" << species_map.node_names[i];
		cout << "\n";
		for (unsigned i = 0; i < species_map.node_names.size(); ++i)
		{
			cout << species_map.node_names[i];
			for (unsigned j = 0; j < species_map.node_names.size(); ++j)
				cout << "\t" << species_map.node_names[species_map.transitions[i][j]];
			cout << "\n";
		}
		cout << species_map.tree << "\n";

	}
	catch (std::exception& e)
	{
		std_print_error(cerr,  e.what());
		return 1;
	}

	return 0;
}


